﻿#include "transferwidget.h"
#include "ui_transferwidget.h"
#include <QFileDialog>
#include <QDebug>
#include "labelinfodlg.h"
#include <QMessageBox>
#include <QProcess>
#include <QTemporaryDir>

TransferWidget::TransferWidget(QWidget *parent) :
    QWidget(parent),
    ui(new Ui::TransferWidget)
{
    ui->setupUi(this);
    this->setWindowTitle(QStringLiteral("YOLOv5模型转换器"));
    m_pluginDir = QCoreApplication::applicationDirPath().replace("Runtimes","TaiJiPlugins");
}

TransferWidget::~TransferWidget()
{
    delete ui;
}

void TransferWidget::on_imgsPath_clicked()
{
    m_imgsPath = QFileDialog::getExistingDirectory(this,QStringLiteral("选择目录"),"./",QFileDialog::ShowDirsOnly);
    if(!m_imgsPath.isEmpty()){
        ui->imgsPath->setText(QStringLiteral("图片路径(√)"));
    }
}

void TransferWidget::on_txtLabelsPath_clicked()
{
    m_txtLabelsPath = QFileDialog::getExistingDirectory(this,QStringLiteral("选择目录"),"./",QFileDialog::ShowDirsOnly);
    if(!m_txtLabelsPath.isEmpty()){
        ui->txtLabelsPath->setText(QStringLiteral("标签路径(√)"));
    }
}

void TransferWidget::on_xmlLabelsPath_clicked()
{
    m_xmlLabelsPath = QFileDialog::getExistingDirectory(this,QStringLiteral("选择目录"),"./",QFileDialog::ShowDirsOnly);
    if(!m_xmlLabelsPath.isEmpty()){
        ui->xmlLabelsPath->setText(QStringLiteral("标签转换路径(√)"));
    }
}

void TransferWidget::on_startlabeltranslate_clicked()
{
    if(!m_imgsPath.isEmpty() && !m_txtLabelsPath.isEmpty() && !m_xmlLabelsPath.isEmpty()){
        LabelInfoDlg dlg;
        if(dlg.exec()==QDialog::Rejected){
            return;
        }
        QTemporaryDir dir;
        if(dir.isValid()){
            QFile::copy( m_pluginDir + "/ModelTransfer/main.py",dir.path()+"/main.py");
            QFile file(dir.path()+"/config.py");
            if (!file.open(QIODevice::WriteOnly|QIODevice::Truncate)) {
                QMessageBox::information(NULL,QStringLiteral("标签转换"),QStringLiteral("无法生成配置文件:")+m_pluginDir+"/ModelTransfer/config.py");
                return;
            }
            //QMessageBox::information(NULL,QStringLiteral("标签转换"),QStringLiteral("完成创建:")+dir.path());
            QString writeinfo;
            writeinfo += "\nIMAGE_DIR = ";
            writeinfo += "\"";
            writeinfo += m_imgsPath;
            writeinfo += "\"";
            writeinfo += "\nLABEL_DIR = ";
            writeinfo += "\"";
            writeinfo += m_txtLabelsPath;
            writeinfo += "\"";
            writeinfo += "\nXML_DIR = ";
            writeinfo += "\"";
            writeinfo += m_xmlLabelsPath;
            writeinfo += "\"";
            writeinfo += "\nnames = [";
            QStringList labels = dlg.m_labelTxt.split("\n");
            foreach(QString label,labels){
                writeinfo += "\""+label + "\",";
            }
            writeinfo += "]";
            file.write(writeinfo.toUtf8());
            file.close();
            QString cmd = QString("python.exe %1 --").arg(dir.path()+"/main.py");
            cmd += ui->labelmethod->currentText().trimmed();
            //system(cmd.toStdString().c_str());
            ui->plainTextEdit->insertPlainText(cmd+"\n");
            QStringList arguments;
            arguments << "/c" << cmd;
            QProcess process(this);
            process.start("cmd.exe", arguments);
            process.waitForStarted();
            process.waitForFinished();

            QString strResult =QString::fromLocal8Bit(process.readAllStandardOutput());
            ui->plainTextEdit->insertPlainText(strResult);
        }else{
            QMessageBox::information(NULL,QStringLiteral("标签转换"),QStringLiteral("临时目录创建失败"));
        }

    }else{
        QMessageBox::information(NULL,QStringLiteral("标签转换"),QStringLiteral("请先设置图片标签路径"));
    }
}

void TransferWidget::on_startmodeltranslate_clicked()
{
    if(!m_weightPath.isEmpty()){
        QString cmd = QString("python.exe %1/ModelTransfer/export.py --weights %2 --include %3  --img %4")
                .arg(m_pluginDir)
                .arg(m_weightPath)
                .arg(ui->exportmodel->currentText().trimmed())
                .arg(960);
        //system(cmd.toStdString().c_str());
        ui->plainTextEdit->insertPlainText(cmd+"\n");
        QStringList arguments;
        arguments << "/c" << cmd;
        QProcess process(this);
        process.start("cmd.exe", arguments);
        process.waitForStarted();
        process.waitForFinished();

        QString strResult =QString::fromLocal8Bit(process.readAllStandardOutput());
        ui->plainTextEdit->insertPlainText(strResult);
    }else{
        QMessageBox::information(NULL,QStringLiteral("模型转换"),QStringLiteral("请先设置模型路径"));
    }
}


void TransferWidget::on_weightpath_clicked()
{
    m_weightPath = QFileDialog::getOpenFileName(NULL,QStringLiteral("选择模型"),"./","*.pt");
    if(!m_weightPath.isEmpty()){
        ui->weightpath->setText(QStringLiteral("模型路径(√)"));
    }
}
